package {{packageName}}.graph;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncEdgeAction;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
{{importSection}}

{{#hasRetriever}}
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
{{/hasRetriever}}

{{#hasCode}}
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionConfig;
import com.alibaba.cloud.ai.graph.node.code.CodeExecutor;
{{/hasCode}}

@Component
public class GraphBuilder {
    {{#hasRetriever}}
    // todo: add your knowledge base
    @Value("${rag.source:classpath:${}")
    private Resource ragSource;
    {{/hasRetriever}}

    @Bean
    public CompiledGraph buildGraph(
        ChatModel chatModel
        {{#hasRetriever}}, VectorStore vectorStore{{/hasRetriever}}
        {{#hasCode}}, CodeExecutionConfig codeExecutionConfig, CodeExecutor codeExecutor{{/hasCode}}
        ) throws GraphStateException {
        ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(new SimpleLoggerAdvisor()).build();

        // new stateGraph
        StateGraph stateGraph = new StateGraph({{stateSection}});
        // add nodes
        {{nodeSection}}
        // add edges
        {{edgeSection}}
        return stateGraph.compile();
    }

    {{#hasRetriever}}
        @Bean
        public VectorStore customVectorStore(EmbeddingModel embeddingModel) {
            var chunks = new TokenTextSplitter()
                .transform(new TextReader(ragSource).read());
            SimpleVectorStore vectorStore = SimpleVectorStore
                .builder(embeddingModel)
                .build();
            vectorStore.write(chunks);
            return vectorStore;
        }
    {{/hasRetriever}}

    {{#hasCode}}
        @Bean
    	public Path tempDir() throws IOException {
    	    // todo： set your work dir
    		Path tempDir = Files.createTempDirectory("code-execution-workdir-");
    		tempDir.toFile().deleteOnExit();
    		return tempDir;
    	}

    	@Bean
    	public CodeExecutionConfig codeExecutionConfig(Path tempDir) {
    		return new CodeExecutionConfig().setWorkDir(tempDir.toString());
    	}

    	@Bean
    	public CodeExecutor codeGenerator() {
    		return new LocalCommandlineCodeExecutor();
    	}
    {{/hasCode}}
}
